package com.tech.distribute_lock.zk;

import com.tech.distribute_lock.DistributeLockStategy;
import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.*;
import org.apache.zookeeper.data.Stat;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.locks.LockSupport;
import java.util.stream.Collectors;
/**
 * 临时顺序节点实现；每个锁，需要单独创建文件夹。方法、or 类型维度，不建议id维度；粒度较粗
 * 临时节点；无法避免惊群
 *
 * 加锁非原子化
 */
@Component
@Slf4j
public class ZookeeperDistributeLockStategy implements DistributeLockStategy {
    private static final String LOCK_ROOT_PATH = "/Locks";
    private static final String LOCK_NODE_NAME = "Lock_";
    /**等待间隔*/
    private static final Long WAIT_STEP_NANAOS = 3*1000*1000*1000L;
    /**lockKey-OwnerId*/
    private static final ConcurrentHashMap<String, String> lockKeyOwnerMap = new ConcurrentHashMap<>();
    /**preLockPath-thread*/
    private static final ConcurrentHashMap<String, Thread> preLockPathThreadMap = new ConcurrentHashMap<>();
    /**lockKey-lockPath*/
    private static final ThreadLocal<ConcurrentHashMap<String, String>> lockKeyLockPathMapTl = ThreadLocal.withInitial(ConcurrentHashMap::new);

    //ZooKeeper配置信息
    private ZooKeeper zkClient;
 
    // 监控lockPath的前一个节点的watcher
    private class PrePathDeletedWatcher implements Watcher{
        @Override
        public void process(WatchedEvent event) {
            if(Event.EventType.NodeDeleted==event.getType()){
                String path = event.getPath();
                log.info(path + " 前锁释放");
                // 监听到前锁释放，唤醒线程
                LockSupport.unpark(preLockPathThreadMap.get(path));
            }
        }
    };

    public static void main(String[] args) throws Exception {
        ZookeeperDistributeLockStategy zookeeperDistributeLockStategy = new ZookeeperDistributeLockStategy();
        Runnable runnable = ()->{
            zookeeperDistributeLockStategy.tryLock("aa", 2000);
            zookeeperDistributeLockStategy.unLock("aa");
        };
        int n = 10;
        ExecutorService executorService = Executors.newFixedThreadPool(n);
        for (int i = 0; i < n; i++) {
            executorService.submit(runnable);
        }
        Thread.sleep(10*1000);
        zookeeperDistributeLockStategy.zkClient.close();
        // 线程池不关闭，进程不终止
        executorService.shutdown();
    }
 
    public ZookeeperDistributeLockStategy() throws IOException {
        zkClient= new ZooKeeper("localhost:2181", 10000, new Watcher() {
            @Override
            public void process(WatchedEvent event) {
                if(event.getState()== Event.KeeperState.Disconnected){
                    log.info("失去连接");
 
                }
            }
        });
    }

    @Override
    public boolean tryLock(String lockKey, int expireTime) {
        try {
            //创建锁节点
            createLock(lockKey);
            //尝试获取锁
            return attemptLock(lockKey);
        } catch (KeeperException e) {
            log.error("", e);
        } catch (InterruptedException e) {
            log.error("", e);
        }
        return false;
    }

    //创建锁的原语实现。在lock节点下创建该线程的锁节点
    private void createLock(String lockKey) throws KeeperException, InterruptedException {
        //如果lockKey节点不存在，则创建
        String lockSubDir = getPath(LOCK_ROOT_PATH, lockKey);
        Stat stat = zkClient.exists(lockSubDir, false);
        if (stat == null) {
            //如果根节点不存在，则创建根节点
            stat = zkClient.exists(LOCK_ROOT_PATH, false);
            if (stat == null) {
                zkClient.create(LOCK_ROOT_PATH, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
            }
            zkClient.create(lockSubDir, new byte[0], ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);
        }

        // 创建EPHEMERAL_SEQUENTIAL类型节点
        String ownerId = getOwnerId();
        String lockPath = zkClient.create(getPath(lockSubDir, LOCK_NODE_NAME),
                ownerId.getBytes(), ZooDefs.Ids.OPEN_ACL_UNSAFE,
                CreateMode.EPHEMERAL_SEQUENTIAL);
        log.info(ownerId + " 锁创建: " + lockPath);
        lockKeyLockPathMapTl.get().put(lockSubDir, lockPath);
    }

    private String getPath(String... paths) {
        return Arrays.stream(paths).collect(Collectors.joining("/"));
    }

    private boolean attemptLock(String lockKey) throws KeeperException, InterruptedException {
        String lockSubDir = getPath(LOCK_ROOT_PATH, lockKey);
        // 获取Lock所有子节点，按照节点序号排序
        List<String> lockNodes = zkClient.getChildren(lockSubDir, false);

        Collections.sort(lockNodes);

        String lockPath = lockKeyLockPathMapTl.get().get(lockSubDir);
        int index = lockNodes.indexOf(lockPath.substring(lockSubDir.length()+1));

        // 如果lockPath是序号最小的节点，则获取锁
        if (index==0) {
            String ownerId = getOwnerId();

            log.info(ownerId + " 锁获得, lockPath: " + lockPath);

            lockKeyOwnerMap.put(lockSubDir, ownerId);
            return true;
        } else if(index>0){
            // lockPath不是序号最小的节点，监控前一个节点
            String preLockPath = getPath(lockSubDir, lockNodes.get(index - 1));

            log.info("监听前锁, lockPath: {}", preLockPath);

            Stat stat = zkClient.exists(preLockPath, new PrePathDeletedWatcher());

            // 假如前一个节点不存在了，比如说执行完毕，或者执行节点掉线，重新获取锁
            if (stat == null) {
                attemptLock(lockKey);
            } else { // 阻塞当前进程，直到preLockPath释放锁，被watcher观察到，notifyAll后，重新acquireLock
                log.info(" 等待前锁释放，prelocakPath："+preLockPath);
                preLockPathThreadMap.put(preLockPath, Thread.currentThread());
                LockSupport.parkNanos(index*WAIT_STEP_NANAOS);
                attemptLock(lockKey);
            }
        }
        return false;
    }

    @Override
    public boolean unLock(String lockKey){
        try {
            String lockSubDir = getPath(LOCK_ROOT_PATH, lockKey);
            String ownerId = lockKeyOwnerMap.get(lockSubDir);
            if(getOwnerId().equals(ownerId)){
                String lockPath = lockKeyLockPathMapTl.get().get(lockSubDir);
                zkClient.delete(lockPath, -1);
//                zkClient.close();
                log.info(" 锁释放：" + lockPath);
                lockKeyLockPathMapTl.get().remove(lockSubDir);
            }
        } catch (InterruptedException e) {
            log.error("", e);
        } catch (KeeperException e) {
            log.error("", e);
        }
        return true;
    }
}